import paramiko


class SSHHandler:
    """SSH连接处理器"""
    
    def __init__(self, host, port, username, password=None, ssh_key=None):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.ssh_key = ssh_key
        self.client = None
        self.channel = None
        self.connected = False
        
    def connect(self):
        """建立SSH连接"""
        try:
            self.client = paramiko.SSHClient()
            self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            connect_kwargs = {
                'hostname': self.host,
                'port': self.port,
                'username': self.username,
                'timeout': 10,
            }
            
            if self.password:
                connect_kwargs['password'] = self.password
            elif self.ssh_key:
                import io
                key_file = io.StringIO(self.ssh_key)
                pkey = paramiko.RSAKey.from_private_key(key_file)
                connect_kwargs['pkey'] = pkey
            
            self.client.connect(**connect_kwargs)
            
            # 打开一个交互式shell
            self.channel = self.client.invoke_shell(term='xterm', width=120, height=30)
            self.channel.settimeout(0.0)
            
            self.connected = True
            return True
            
        except Exception as e:
            self.connected = False
            raise Exception(f"SSH连接失败: {str(e)}")
    
    def send(self, data):
        """发送数据到SSH通道"""
        if self.channel and self.connected:
            try:
                self.channel.send(data)
            except Exception as e:
                self.connected = False
                raise Exception(f"发送数据失败: {str(e)}")
    
    def recv(self, size=1024):
        """从SSH通道接收数据"""
        if self.channel and self.connected:
            try:
                if self.channel.recv_ready():
                    return self.channel.recv(size)
                return b''
            except Exception as e:
                self.connected = False
                return b''
        return b''
    
    def resize(self, width, height):
        """调整终端大小"""
        if self.channel and self.connected:
            try:
                self.channel.resize_pty(width=width, height=height)
            except Exception as e:
                pass
    
    def close(self):
        """关闭SSH连接"""
        self.connected = False
        if self.channel:
            try:
                self.channel.close()
            except:
                pass
        if self.client:
            try:
                self.client.close()
            except:
                pass
    
    def is_active(self):
        """检查连接是否活跃"""
        return self.connected and self.channel and not self.channel.closed
